# -*- coding: utf-8 -*-

"""
T-M模型

"""

import numpy as np

from scipy.optimize import leastsq

from db import get_data_by_proc
from util import get_r_squared


def model(p, x):
    """T-M model:Rp - Rf = Alpha + Beta1 * (Rm - Rf) + Beta2 * (Rm - Rf) ^ 2"""
    alpha, beta1, beta2 = p
    return alpha + beta1 * x + beta2 * x ** 2


def residual(p, x, y):
    """残差"""
    return y - model(p, x)

def get_t_m():

    fund_id = 'MF00003PWJ'
    current_month = '2020-02'
    month_count = 12

    # 返回结果
    result_dict = {}

    proc_name = 'sp_get_fund_ret_py'
    params = (fund_id, current_month, month_count)
    columns = ['A', 'Rp', 'Rm', 'Rf']
    fund_df = get_data_by_proc(proc_name=proc_name, params=params, columns=columns)
    fund_df.drop(columns=['A'], inplace=True)
    x0 = np.array(fund_df['Rm'])
    y0 = np.array(fund_df['Rp'])

    # 初始值
    p0 = np.random.randn(3)

    # 调用leastsq进行数据拟合
    p = leastsq(func=residual, x0=p0, args=(x0, y0))
    # 拟合结果
    alpha, beta1, beta2 = p[0]
    result_dict['alpha'] = alpha
    result_dict['beta1'] = beta1
    result_dict['beta2'] = beta2

    # 计算模拟
    y_test = np.array(model(p[0], x0))
    y_true = y0

    # 计算拟合度r_squared
    r_squared = get_r_squared(y_true, y_test)
    result_dict['r_squared'] = r_squared
    return result_dict


if __name__ == '__main__':

    print(get_t_m())

    pass
